/*
 * Copyright 2021 TiKV Project Authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

package com.webank.tikv.util;

import com.google.common.collect.ImmutableList;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import io.grpc.ManagedChannel;
import io.grpc.netty.NettyChannelBuilder;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslContextBuilder;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

import javax.net.ssl.SSLException;
import java.io.File;
import java.net.URI;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class ChannelFactory implements AutoCloseable {

    private static final Logger LOGGER = LogManager.getLogger(ChannelFactory.class);

    private final int maxFrameSize;
    private final int keepAliveTime;
    private final int keepAliveTimeout;
    private final int idleTimeout;

    // After `connRecycleTime` seconds elapses, the old channels will be forced to shut down,
    // to avoid using the old context all the time including potential channel leak.
    private long connRecycleTime;

    private CertContext certContext;
    private CertWatcher certWatcher;

    private ScheduledExecutorService recycler;

    private final ConcurrentHashMap<String, ManagedChannel> connPool = new ConcurrentHashMap<>();
    private final AtomicReference<SslContextBuilder> sslContextBuilder = new AtomicReference<>();
    private final ReadWriteLock lock = new ReentrantReadWriteLock();

    public class CertWatcher {

        private final List<File> targets;
        private final List<Long> lastReload = new ArrayList<>();
        private final Runnable onChange;

        public CertWatcher(long pollInterval, List<File> targets, Runnable onChange) {
            this.targets = targets;
            this.onChange = onChange;

            for (File ignored : targets) {
                lastReload.add(0L);
            }

            recycler.scheduleAtFixedRate(this::tryReload, pollInterval, pollInterval, TimeUnit.MILLISECONDS);
        }

        // If any execution of the task encounters an exception, subsequent executions are suppressed.
        private void tryReload() {
            // Add exception handling to avoid schedule stop.
            try {
                if (needReload()) {
                    onChange.run();
                }
            } catch (Exception e) {
                LOGGER.error("Failed to reload cert", e);
            }
        }

        private boolean needReload() {
            boolean needReload = false;
            // Check all the modification of the `targets`.
            // If one of them changed, means to need reload.
            for (int i = 0; i < targets.size(); i++) {
                try {
                    long lastModified = targets.get(i).lastModified();
                    if (lastModified != lastReload.get(i)) {
                        lastReload.set(i, lastModified);
                        LOGGER.warn("detected ssl context changes: {}", targets.get(i));
                        needReload = true;
                    }
                } catch (Exception e) {
                    LOGGER.error("fail to check the status of ssl context files", e);
                }
            }
            return needReload;
        }

    }

    public ChannelFactory(int maxFrameSize, int keepAliveTime, int keepAliveTimeout, int idleTimeout, long connRecycleTime) {
        this.maxFrameSize = maxFrameSize;
        this.keepAliveTime = keepAliveTime;
        this.keepAliveTimeout = keepAliveTimeout;
        this.idleTimeout = idleTimeout;
        this.connRecycleTime = connRecycleTime;

        this.recycler = Executors.newSingleThreadScheduledExecutor(
            new ThreadFactoryBuilder().setNameFormat("tikv-ch-%d").setDaemon(true).build());
        recycler.scheduleAtFixedRate(this::evictInvalidChannel, connRecycleTime, connRecycleTime, TimeUnit.MILLISECONDS);
    }

    public ChannelFactory(int maxFrameSize, int keepAliveTime, int keepAliveTimeout, int idleTimeout, long connRecycleTime,
                          long certReloadInterval, String trustCertCollectionFilePath, String keyCertChainFilePath, String keyFilePath) {
        this(maxFrameSize, keepAliveTime, keepAliveTimeout, idleTimeout, connRecycleTime);

        this.certContext = new CertContext.OpenSslContext(trustCertCollectionFilePath, keyCertChainFilePath, keyFilePath);
        File trustCert = new File(trustCertCollectionFilePath);
        File keyCert = new File(keyCertChainFilePath);
        File key = new File(keyFilePath);

        if (certReloadInterval > 0) {
            onCertChange();
            this.certWatcher = new CertWatcher(certReloadInterval, ImmutableList.of(trustCert, keyCert, key), this::onCertChange);
        }
    }

    public ChannelFactory(int maxFrameSize, int keepAliveTime, int keepAliveTimeout, int idleTimeout, long connRecycleTime,
                          long certReloadInterval, String jksKeyPath, String jksKeyPassword, String jksTrustPath, String jksTrustPassword) {
        this(maxFrameSize, keepAliveTime, keepAliveTimeout, idleTimeout, connRecycleTime);

        this.certContext = new CertContext.JksContext(jksKeyPath, jksKeyPassword, jksTrustPath, jksTrustPassword);
        File jksKey = new File(jksKeyPath);
        File jksTrust = new File(jksTrustPath);

        if (certReloadInterval > 0) {
            onCertChange();
            this.certWatcher = new CertWatcher(certReloadInterval, ImmutableList.of(jksKey, jksTrust), this::onCertChange);
        }
    }

    private void onCertChange() {
        try {
            SslContextBuilder newBuilder = certContext.createSslContextBuilder();
            lock.writeLock().lock();
            sslContextBuilder.set(newBuilder);

            List<ManagedChannel> pending = new ArrayList<>(connPool.values());
            recycler.schedule(() -> cleanExpiredConn(pending), connRecycleTime, TimeUnit.MILLISECONDS);

            connPool.clear();
        } finally {
            lock.writeLock().unlock();
        }
    }

    private void evictInvalidChannel() {
        Iterator<Map.Entry<String, ManagedChannel>> iterator = connPool.entrySet().iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, ManagedChannel> entry = iterator.next();
            if (entry.getValue().isShutdown()) {
                iterator.remove();
                LOGGER.warn("evict invalid channel:{}", entry.getKey());
            }
        }
    }

    public ManagedChannel getChannel(String address) {
        if (certContext != null) {
            try {
                lock.readLock().lock();
                return connPool.computeIfAbsent(
                    address, key -> createChannel(sslContextBuilder.get(), address));
            } finally {
                lock.readLock().unlock();
            }
        }
        return connPool.computeIfAbsent(address, key -> createChannel(null, address));
    }

    private ManagedChannel createChannel(SslContextBuilder sslContextBuilder, String address) {
        URI uri;
        try {
            uri = UriUtils.addrToUri(address);
        } catch (Exception e) {
            throw new IllegalArgumentException("failed to form address " + address, e);
        }

        // Channel should be lazy without actual connection until first call
        // So a coarse grain lock is ok here
        NettyChannelBuilder builder = NettyChannelBuilder.forAddress(uri.getHost(), uri.getPort())
            .maxInboundMessageSize(maxFrameSize)
            .keepAliveTime(keepAliveTime, TimeUnit.MILLISECONDS)
            .keepAliveTimeout(keepAliveTimeout, TimeUnit.MILLISECONDS)
            .keepAliveWithoutCalls(true)
            .idleTimeout(idleTimeout, TimeUnit.MILLISECONDS);

        if (sslContextBuilder == null) {
            return builder.usePlaintext().build();
        } else {
            SslContext sslContext;
            try {
                sslContext = sslContextBuilder.build();
            } catch (SSLException e) {
                LOGGER.error("create ssl context failed: {}", e.getMessage());
                throw new IllegalArgumentException(e);
            }
            return builder.sslContext(sslContext).build();
        }
    }

    private void cleanExpiredConn(List<ManagedChannel> pending) {
        LOGGER.info("cleaning expire channels");
        for (ManagedChannel channel : pending) {
            channel.shutdownNow();
            while (!channel.isShutdown()) {
                try {
                    channel.awaitTermination(5, TimeUnit.SECONDS);
                } catch (Exception e) {
                    LOGGER.warn("recycle channels timeout: {}", e.getMessage());
                }
            }
        }
    }

    public void close() {
        for (ManagedChannel ch : connPool.values()) {
            ch.shutdown();
        }
        connPool.clear();
        if (recycler != null) {
            recycler.shutdown();
        }
    }
}
